package chatBot.service.impl;

import chatBot.service.TokenConsumptionService;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import chatBot.event.GptEventSourceListener;
import chatBot.exception.NoPermissionException;
import chatBot.exception.NoTokenCountException;
import chatBot.mapper.MessageMapper;
import chatBot.mapper.ConversationMapper;
import chatBot.mapper.UserMapper;
import chatBot.entity.Message;
import chatBot.entity.Conversation;
import chatBot.entity.User;
import chatBot.service.ChatgptService;
import chatBot.utils.ContextHolderUtil;
import com.plexpt.chatgpt.ChatGPTStream;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import javax.annotation.Resource;
import java.io.IOException;
import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.ArrayList;


@Slf4j
@Service
public class ChatgptServiceImpl extends ServiceImpl<MessageMapper, Message> implements ChatgptService {

    @Resource
    UserMapper userMapper;

    @Resource
    ConversationMapper conversationMapper;

    @Value("${my-conf.gpt-key}")
    private String apiKey;
    @Value("${my-conf.apiURl}")
    private String apiURl;

    @Autowired
    private TokenConsumptionService tokenConsumptionService;

    /**
     * 代表消息发送方向
     */
    private static final int DIRECTION_QUESTION = 0;
    private static final int DIRECTION_ANSWER = 1;

    /**
     * 发送消息时加入的上下文信息数
     */
    private static final int CONTEXT_MESSAGE_NUM = 2;

    @Override
    public List<Message> messagesOfConversation(Long conversationId, Integer showItem) {
        // 1、校验会话是否属于当前登录用户
        Long userId = ContextHolderUtil.getUserId();
        Conversation conversation = conversationMapper.selectById(conversationId);
        if (userId == null) {
            throw new NoPermissionException("用户未登录");
        }
        if (!userId.equals(conversation.getUserId())) {
            throw new NoPermissionException("不能查询其他用户的会话！");
        }
        // 2、查询conversation的最新showItem条会话
        LambdaQueryWrapper<Message> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        lambdaQueryWrapper.eq(Message::getConversationId, conversationId);
        lambdaQueryWrapper.orderByAsc(Message::getMessageId);
        lambdaQueryWrapper.last(showItem != null, "limit " + showItem);
        return list(lambdaQueryWrapper);
    }

    /**
     * 将本项目的数据类型Message转为SDK需要的Message类格式，可以Ctrl+戳戳查看Message类型
     * @param src 数据库格式消息列表
     * @return SDK所需格式消息列表
     */
    private List<com.plexpt.chatgpt.entity.chat.Message> transform(List<Message> src) {
        List<com.plexpt.chatgpt.entity.chat.Message> tgt = new ArrayList<>();
        // 对src列表中的每条消息
        for (Message message: src) {
            // 根据消息类型调用对应的方法，向tgt中加入所需类型消息
            tgt.add(message.getMessageDirection() == DIRECTION_QUESTION ?
                    // 用户发送的消息
                    com.plexpt.chatgpt.entity.chat.Message.of(message.getContent())
                    // GPT发送的消息
                    : com.plexpt.chatgpt.entity.chat.Message.ofAssistant(message.getContent()));
        }
        return tgt;
    }


    @Override
    public SseEmitter sendMsgSse(Long conversationId, String prompt, boolean isNewConversation) {
        // 获取SecurityContextHolder中的用户id, 判断该会话是否属于当前用户
        Long userId = ContextHolderUtil.getUserId();
        Conversation conversation = conversationMapper.selectById(conversationId);
        if (userId == null) {
            throw new NoPermissionException("用户未登录");
        }
        if (!userId.equals(conversation.getUserId())) {
            throw new NoPermissionException("不能更改其他用户的会话！");
        }

        // 得到用户
        User user = userMapper.selectById(userId);

        // 计算Prompt的Token数量
        int promptTokens = tokenConsumptionService.calculateTokenCount(prompt);

        // 计算上下文Token数量
        int contextTokens = conversationId != null
                ? messagesOfConversation(conversationId, CONTEXT_MESSAGE_NUM).stream()
                .mapToInt(msg -> tokenConsumptionService.calculateTokenCount(msg.getContent())).sum()
                : 0;

        // 估算最大响应Token数量
        int maxResponseTokens = tokenConsumptionService.calculateMaxResponseTokens(promptTokens, contextTokens);

        // 计算总Token消耗
        int totalConsumption = tokenConsumptionService.calculateTotalTokenConsumption(promptTokens, contextTokens, maxResponseTokens);

        // 检查用户是否有足够的Token
        if (user.getUserTokenCount() < totalConsumption) {
            throw new NoTokenCountException("用户余额不足");
        }

        // 实例化ChatGPTStream工具类对象
        ChatGPTStream chatgptStream = ChatGPTStream.builder()
                .timeout(50)
                .apiKey(apiKey) // 绑定API Key
                .apiHost(apiURl)
                .build()
                .init();

        // 实例化流式输出类
        SseEmitter sseEmitter = new SseEmitter(-1L);
        GptEventSourceListener listener = new GptEventSourceListener(sseEmitter);

        // 加入历史消息记录，提供上下文信息
        List<com.plexpt.chatgpt.entity.chat.Message> messages = conversationId != null
                ? transform(messagesOfConversation(conversationId, CONTEXT_MESSAGE_NUM))
                : new ArrayList<>();

        // 加入本次提问问题
        com.plexpt.chatgpt.entity.chat.Message message = com.plexpt.chatgpt.entity.chat.Message.of(prompt);
        messages.add(message);
        log.info("messages: {}", messages);

        // 设置完成时的回调函数
        listener.setOnComplete(msg -> {
            // 计算实际响应的Token数量
            int responseTokens = tokenConsumptionService.calculateTokenCount(msg);

            // 更新总Token消耗
            int finalConsumption = tokenConsumptionService.calculateTotalTokenConsumption(promptTokens, contextTokens, responseTokens);

            // 扣除用户的token并更新账户信息
            user.setUserTokenCount(user.getUserTokenCount() - finalConsumption);
            userMapper.updateById(user);

            // 保存历史信息
            saveMessage(conversationId, prompt, DIRECTION_QUESTION);
            saveMessage(conversationId, msg, DIRECTION_ANSWER);

            try {
                // 仅在新会话时发送结束标志消息
                if (isNewConversation) {
                    Map<String, Object> endMessage = new HashMap<>();
                    endMessage.put("conversationId", conversationId);
                    endMessage.put("userTokenCount", user.getUserTokenCount());
                    String conversationName = conversationMapper.selectById(conversationId).getConversationName();
                    endMessage.put("conversationName", conversationName); // 添加会话名称
                    endMessage.put("endOfStream", true);
                    sseEmitter.send(SseEmitter.event().data(endMessage));
                }
            } catch (IOException e) {
                e.printStackTrace();
            } finally {
                sseEmitter.complete();
            }
        });

        // 提问
        chatgptStream.streamChatCompletion(messages, listener);

        return sseEmitter;
    }


    /**
     * 将消息保存至数据库中
     * @param conversationId 会话id
     * @param msg 消息内容
     * @param messageDirection 消息方向
     */
    public void saveMessage(Long conversationId, String msg, Integer messageDirection) {
        Message message = new Message();
        message.setContent(msg);
        message.setMessageDirection(messageDirection);
        message.setConversationId(conversationId);
        message.setCreateTime(LocalDateTime.now());
        save(message);
    }

}
